import os
import sys
sys.path.append(os.path.abspath('.'))

from xml.dom.minidom import parse
import cv2
import matplotlib.pyplot as plt
import numpy as np

import csv

'''
result
[367.8 402.8]
'''

f = open('./scripts/img_size.csv','w',encoding='utf-8')
csv_writer = csv.writer(f)

root_img_path='/home/hfle/data_dog/low-resolution'
root_label_path='/home/hfle/data_dog/Low-Annotations'
class_names=os.listdir(root_img_path)

total=list()
csv_writer.writerow(['No','img_class','num','img_height','img_width','head_height','head_width','body_height','body_width'])
for i,c in enumerate(class_names):
    img_dir=os.path.join(root_img_path,c)

    attr_list=list()
    for img_name in os.listdir(img_dir):
        img_path=os.path.join(img_dir,img_name)
        label_path=os.path.join(root_label_path,c,img_name)+'.xml'

        if not (os.path.exists(img_path) and os.path.exists(label_path)):
            continue

        # 读取图片
        img = cv2.imread(img_path)
        img_height,img_width=img.shape[:2]

        # 读取xml
        DOMTree = parse(label_path)
        collection = DOMTree.documentElement
        headbndbox = collection.getElementsByTagName("headbndbox")[0]
        hxmin = int(headbndbox.getElementsByTagName("xmin")[0].childNodes[0].data)
        hymin = int(headbndbox.getElementsByTagName("ymin")[0].childNodes[0].data)
        hxmax = int(headbndbox.getElementsByTagName("xmax")[0].childNodes[0].data)
        hymax = int(headbndbox.getElementsByTagName("ymax")[0].childNodes[0].data)

        bodybndbox = collection.getElementsByTagName("bodybndbox")[0]
        bxmin = int(bodybndbox.getElementsByTagName("xmin")[0].childNodes[0].data)
        bymin = int(bodybndbox.getElementsByTagName("ymin")[0].childNodes[0].data)
        bxmax = int(bodybndbox.getElementsByTagName("xmax")[0].childNodes[0].data)
        bymax = int(bodybndbox.getElementsByTagName("ymax")[0].childNodes[0].data)

        # 框
        head_height=hymax-hymin
        head_width=hxmax-hxmin
        head_x=(hxmin+hxmax)/2.0
        head_y=(hymin+hymax)/2.0

        body_height=bymax-bymin
        body_width=bxmax-bxmin
        body_x=(bxmin+bxmax)/2.0
        body_y=(bymin+bymax)/2.0

        attr_list.append([img_height,img_width,head_height,head_width,body_height,body_width])
    total.append(np.around(np.array(attr_list).mean(axis=0),1))
    print(c,len(os.listdir(img_dir)),total[-1])
    csv_writer.writerow([str(i+1),c,len(os.listdir(img_dir))]+list(total[-1]))
    f.flush()
res=np.array(total)
avg_res=res.mean(axis=0)
print(avg_res)


